import pandas as pd
import numpy as np

def add_features(df):
    """添加特征工程"""
    # 示例特征：计算奇偶比、质数个数等
    df['odd_even_ratio'] = df['numbers'].apply(lambda x: sum([1 for n in x if n % 2 == 1]) / len(x))
    df['prime_count'] = df['numbers'].apply(lambda x: sum([1 for n in x if is_prime(n)]))
    return df

def is_prime(num):
    """判断质数的辅助函数"""
    if num < 2:
        return False
    for i in range(2, int(np.sqrt(num)) + 1):
        if num % i == 0:
            return False
    return True

def save_features(df, file_path):
    """保存特征数据"""
    df.to_csv(file_path, index=False)

if __name__ == "__main__":
    processed_data_path = "data/processed/ssq/cleaned_data.csv"
    features_data_path = "data/processed/ssq/features_data.csv"

    df = pd.read_csv(processed_data_path)
    df = add_features(df)
    save_features(df, features_data_path)